package com.whz.governance.redislimiter.tool;

import com.whz.governance.redislimiter.command.ScriptParams;
import com.whz.governance.redislimiter.command.TokenBucketCommand;
import com.whz.governance.redislimiter.command.TokenBucketResult;
import lombok.experimental.UtilityClass;
import org.springframework.util.Assert;

import java.time.temporal.ChronoUnit;
import java.util.List;

/**
 * @author whz
 */
@UtilityClass
public class TokenBucketUtils {
    private static final TokenBucketCommand COMMAND = new TokenBucketCommand();

    /**
     * 时间单位为秒的令牌桶执行命令
     *
     * @param key             关键字
     * @param replenishRate   令牌放入速度(单位秒)
     * @param burstCapacity   令牌桶容量
     * @param requestedTokens 每个请求token数,一般是1
     * @return 执行结果
     */
    public TokenBucketResult execute(String key, int replenishRate, int burstCapacity, int requestedTokens) {
        return doExecute(key, replenishRate, burstCapacity, requestedTokens, ChronoUnit.SECONDS);
    }

    private TokenBucketResult doExecute(String key, int replenishRate, int burstCapacity, int requestedTokens,
                                        ChronoUnit unit) {
        Assert.hasText(key, "TokenBucket: key not null");
        Assert.isTrue(replenishRate > 0, "TokenBucket: replenishRate > 0");
        Assert.isTrue(burstCapacity > replenishRate, "TokenBucket:burstCapacity > replenishRate * unit(to Seconds)");
        Assert.isTrue(requestedTokens > 0, "TokenBucket:requestedTokens > 0");
        ScriptParams params = ScriptParams.create()
                                          .addKey(key + "_token")
                                          .addKey(key + "_timestamp")
                                          .addArg(String.valueOf(replenishRate))
                                          .addArg(String.valueOf(burstCapacity))
                                          .addArg(InstantUtils.currentUnitDuration(unit) + "")
                                          .addArg(String.valueOf(requestedTokens))
                                          .addArg(InstantUtils.getTtl(replenishRate, burstCapacity, unit) * 2 + "");
        List<Long> execute = COMMAND.execute(params);
        TokenBucketResult result = new TokenBucketResult();
        result.setPermit(execute.get(0) == 1);
        result.setTokensLeft(execute.get(1));
        return result;
    }


    /**
     * 令牌桶执行命令
     *
     * @param key             关键字
     * @param replenishRate   令牌放入速度
     * @param burstCapacity   令牌桶容量
     * @param requestedTokens 每个请求token数,一般是1
     * @param unit            时间单位
     * @return 执行结果
     */
    public TokenBucketResult execute(String key, int replenishRate, int burstCapacity, int requestedTokens,
                                     ChronoUnit unit) {
        return doExecute(key, replenishRate, burstCapacity, requestedTokens, unit);
    }
}